import numpy as np
import matplotlib
matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# ====================== 1. 定义信号 ======================
t = np.linspace(0, 1, 1000, endpoint=False)
f = np.sign(np.sin(2 * np.pi * 3 * t)) + 0.5 * np.cos(2 * np.pi * 7 * t)

# ====================== 2. 傅里叶系数 ======================
N_harmonics = 20  # 展示20个谐波
a = np.zeros(N_harmonics + 1)
b = np.zeros(N_harmonics + 1)

for n in range(N_harmonics + 1):
    a[n] = 2 * np.sum(f * np.cos(2 * np.pi * n * t)) / len(t)
    b[n] = 2 * np.sum(f * np.sin(2 * np.pi * n * t)) / len(t)

amp = np.sqrt(a**2 + b**2)  # 幅值谱
freq = np.arange(N_harmonics + 1)

# ====================== 3. 创建画布（上下结构） ======================
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
fig.tight_layout(pad=3)

# 上半部分：原信号 + 重构信号
ax1.plot(t, f, label='Original Signal', color='gray', alpha=0.5)
reconstructed = np.zeros_like(f)
line1, = ax1.plot(t, reconstructed, color='blue', label='Reconstructed')
line2, = ax1.plot(t, np.zeros_like(f), color='red', label='Current Harmonic')
ax1.set_title("累加重构信号")
ax1.legend()
ax1.grid(True)

# 下半部分：频谱图
bars = ax2.bar(freq, amp, color='skyblue')
ax2.set_title("频谱幅值图")
ax2.set_xlabel("频率 n")
ax2.set_ylabel("幅值")
ax2.grid(True)

# ====================== 4. 动画更新函数 ======================
def update(n):
    global reconstructed
    if n == 0:
        reconstructed = np.ones_like(f) * (a[0] / 2)
        current = reconstructed.copy()
    else:
        harmonic = a[n] * np.cos(2 * np.pi * n * t) + b[n] * np.sin(2 * np.pi * n * t)
        reconstructed += harmonic
        current = harmonic

    # 更新上半部分
    line1.set_ydata(reconstructed)
    line2.set_ydata(current)
    ax1.set_title(f"累加到第 n={n} 个谐波")

    # 更新频谱柱高亮当前频率
    for i, bar in enumerate(bars):
        if i == n:
            bar.set_color('red')
        else:
            bar.set_color('skyblue')

    return line1, line2, *bars

# ====================== 5. 创建动画 ======================
# blit=False 是 matplotlib.animation.FuncAnimation 里的一个参数，它决定 动画每一帧是“局部刷新”还是“整张图重绘”
ani = FuncAnimation(fig, update, frames=N_harmonics + 1, interval=2000, blit=False, repeat=False)

plt.show()
